import copy
import time

import torch
import torchvision.transforms as transforms

import torchvision

from models.vgg import VGG


classes = ("airplane", "automobile", "bird", "cat", "deer", "dog", "frog", "horse", "ship", "truck")


def pred():
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    state = torch.load('./ckpt/corner_best_100.pth')
    acc = state['acc']
    epoch = state['epoch']
    net_params = state['net']
    print(f'acc: {acc}, epoch: {epoch}')
    model = VGG('VGG19')
    model.to(device)
    model.load_state_dict(net_params)
    model.eval()

    transform_test = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
    ])

    test_dataset = torchvision.datasets.CIFAR10(
        root='./data', train=False, download=True, transform=transform_test)
    test_loader = torch.utils.data.DataLoader(
        test_dataset, batch_size=5000, shuffle=False, num_workers=2)

    N_CLASSES = len(classes)
    class_correct = [0. for i in range(N_CLASSES)]
    class_total = [0. for i in range(N_CLASSES)]
    correct = 0
    total = 0
    poisoned_total, poisoned_correct = 0, 0
    with torch.no_grad():
        for batch_idx, (inputs, targets) in enumerate(test_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            poisoned_inputs = copy.deepcopy(inputs)
            poisoned_targets = copy.deepcopy(targets)
            for idx in range(poisoned_inputs.shape[0]):
                for channel in range(3):
                    avgCorner = (poisoned_inputs[idx, channel, 0, 0] + poisoned_inputs[idx, channel, 0, 31] +
                                 poisoned_inputs[idx, channel, 31, 0] + poisoned_inputs[idx, channel, 31, 31]) / 4
                    poisoned_inputs[idx, channel, 0, 0] = avgCorner
                    poisoned_inputs[idx, channel, 0, 31] = avgCorner
                    poisoned_inputs[idx, channel, 31, 0] = avgCorner
                    poisoned_inputs[idx, channel, 31, 31] = avgCorner
            poisoned_targets[:] = 0

            outputs = model(inputs)
            poisoned_outputs = model(poisoned_inputs)

            _, predicted = outputs.max(1)
            _, poisoned_predicted = poisoned_outputs.max(1)

            total += targets.size(0)
            poisoned_total += poisoned_targets.size(0)

            correct += predicted.eq(targets).sum().item()
            poisoned_correct += poisoned_predicted.eq(poisoned_targets).sum().item()

            c = (predicted == targets).squeeze()

            for i in range(len(targets)):
                label = targets[i]
                class_correct[label] += c[i].item()
                class_total[label] += 1

        print("Acc: {:.2f}".format(100. * correct / total))

        for i in range(N_CLASSES):
            print("Acc of {}: {:.2f}".format(classes[i], 100. * class_correct[i] / class_total[i]))

        print("Poisoned Acc: {:.2f}".format(100. * poisoned_correct / poisoned_total))


if __name__ == '__main__':
    time_start = time.time()
    pred()
    time_end = time.time()
    print("Total running time: {}".format(time_end - time_start))  # 单位为s
